# encoding=utf-8
import logging
import time

import cv2
import numpy
from PIL import Image
from TrtInfer import TRTInferencer

model = TRTInferencer("unet-engine.trt", max_batch_size=1)
logging.basicConfig(level=logging.DEBUG)

def pre(data_in: str):
    image = Image.open(data_in)
    _image = image.resize((512, 512), Image.BICUBIC)
    np_image = numpy.array(_image)
    _np_image = numpy.transpose(np_image, (2, 0, 1))
    __np_image = numpy.expand_dims(_np_image, axis=0)
    return __np_image.astype(numpy.float32) / 255


def post(data_in):
    y = numpy.squeeze(data_in) * 255.0

    # 下面照抄
    predict_list = []
    for index in range(y.shape[0]):
        predict = Image.fromarray(y[index].squeeze().astype(numpy.uint8), mode="L")  # 灰度模式存图，h * w无通道
        predict_resize = numpy.array(predict.resize((2816, 1880)))
        predict_list.append(predict_resize)

    # {"background": 0, "hemorrhages": 1, "hard_exudates": 2, "microaneurysms": 3, "disc": 4, "soft_exudates": 5}
    predict_list = predict_list[0], predict_list[2], predict_list[3], predict_list[1], predict_list[5], predict_list[4]
    prob_arr = numpy.array(predict_list) / 255.0
    prob_arr /= numpy.sum(prob_arr, axis=0)
    pred_arr = numpy.zeros_like(prob_arr).astype(numpy.uint8)
    pred_label_arr = numpy.argmax(prob_arr, axis=0)

    for j in range(pred_arr.shape[0]):
        pred_arr[j, :, :][pred_label_arr == j] = 255

    return pred_arr


def show(pred_arr: numpy.ndarray):
    obj_lst = ["background", "hemorrhages", "hard_exudates", "microaneurysms", "disc", "soft_exudates"]
    seg_color_list = [(0, 0, 0), (0, 0, 255), (0, 255, 0), (0, 255, 255), (255, 0, 0), (255, 0, 255)]

    pred_flag_arr = numpy.argmax(pred_arr, axis=0)
    pred_color_arr = numpy.zeros((pred_flag_arr.shape[0], pred_flag_arr.shape[1], 3), numpy.uint8)
    for i in range(1, len(obj_lst)):
        pred_color_arr[pred_flag_arr == i] = seg_color_list[i]

    cv2.imwrite("out.png", pred_color_arr)


if __name__ == '__main__':
    data_in = "test_image.jpeg"
    for i in range(3):
        _s1 = time.perf_counter()
        pre_data = pre(data_in)
        _s2 = time.perf_counter()
        out = model.infer_batch(pre_data)
        post_data_in = out[0].reshape(1, 6, 512, 512)
        _s3 = time.perf_counter()
        post_data = post(post_data_in)
        _s4 = time.perf_counter()
        show(post_data)
        print("loop time: ", i)
        print("\t pre cost: ", _s2 - _s1)
        print("\t trt cost: ", _s3 - _s2)
        print("\t post cost: ", _s4 - _s3)
        print("\t rate ", _s4 - _s2)
        print("\t all cost: ", _s4 - _s1)